package matchbox;

import java.util.ArrayList;

import probability.GaussianDistribution;
import probability.Indicator;
import probability.ProbabilityDistribution;
import utils.Printer;
import basics.DenseMatrix;
import basics.DenseVector;
import basics.Vector;
import basics.VectorMatrix;
import inference.GraphEdge;
import inference.GraphFactor;
import inference.GraphNode;
import inference.GraphUtils;
import inference.GraphVariable;
import inference.BeliefMessage;

public class ZBRFactor extends MatchboxFactor {

	@Override
	public BeliefMessage toCached(GraphNode next) {
		BeliefMessage m = null;
		if (next.is(NodeNames.RT)) {
			m = new BeliefMessage(_toR);
		} else if (next.is(NodeNames.Z)) {
			int idx = GraphUtils.getIndexInEdge(edges(NodeNames.Z), next);
			m = new BeliefMessage(_toz.get(idx));
		} else if (next.is(NodeNames.B)) {
			m = new BeliefMessage(_toz.get(_toz.size() - 1));
		}

		return m;
	}

	private GaussianDistribution _toR = new GaussianDistribution(0, 1);

	// this value will be cached from the down coming believes from z to r
	private ArrayList<GaussianDistribution> _zdistributions = new ArrayList<GaussianDistribution>();

	// this is the distribution cached from the believes going up from r to z
	private ArrayList<GaussianDistribution> _toz = new ArrayList<GaussianDistribution>();

	private boolean _forward = true;

	private GaussianDistribution _fromB;

	public ZBRFactor(String vname) {
		super(vname, Type.Operation);
	}

	@Override
	public BeliefMessage toCached(GraphNode next, String opt) {
		BeliefMessage m = null;
		if (next.is(NodeNames.RT)) {
			m = new BeliefMessage(_toR);
		} else if (next.is(NodeNames.Z)) {
			int idx = GraphUtils.getIndexInEdge(edges(NodeNames.Z), next);
			m = new BeliefMessage(_toz.get(idx));
		} else if (next.is(NodeNames.B)) {
			m = new BeliefMessage(_toz.get(_toz.size() - 1));
		}

		return m;
	}

	@Override
	public GaussianDistribution ep(GraphNode next) {
		GaussianDistribution g = super.ep(next);
		if (next.is(NodeNames.RT)) {
			_toR = g;
		} else if (next.is(NodeNames.Z)) {
			int idx = GraphUtils.getIndexInEdge(edges(NodeNames.Z), next);
			_toz.set(idx, g);
		}

		return g;
	}

	@Override
	public BeliefMessage to(GraphNode next, String opt) {
		BeliefMessage m = null;
		if (next.is(NodeNames.RT)) {
			_forward = true;
			_zdistributions.clear();

			for (GraphNode n : edges(NodeNames.Z).nodes()) {
				m = n.to(this, opt);
				_zdistributions.add((GaussianDistribution) m.getDistribution());
				if (_toz.size() < _zdistributions.size()) {
					_toz.add(new GaussianDistribution(1, 1000));
				}
			}

			// _fromB = (GaussianDistribution) edges(NodeNames.B).nodes().get(0).to(this).getDistribution();
			// _zdistributions.add(_fromB);
			// if (_toz.size() < _zdistributions.size()) {
			// _toz.add(new GaussianDistribution(1, 1000));
			// }
			// _toz.add(new GaussianDistribution(0, 1));

			// _zdistributions.add(new GaussianDistribution(_imp_mean, .1));
			// _toz.add(new GaussianDistribution(_imp_mean, .1));

			// _gtoR = GaussianDistribution.product(_zdistributions);

			_toR = Indicator.trueSkillForwardEqualSum(_zdistributions, _toz);
			_toR.setSigma2(_toR.getSigma2() * .2);
			// _toR.setMu(_toR.getMu() + _imp_mean);

			m = new BeliefMessage(_toR);
		} else if (next.is(NodeNames.Z) || next.is(NodeNames.B)) {
			// The mean and variance of each edge is distributed by subtracting the summation of others from the
			// incoming message.
			// PS.: I think there is an assumption here that the changes in the distributions are small
			// GaussianDistribution tozDist = (GaussianDistribution) edges(NodeNames.RT).nodes().get(0).to(this)
			// .getDistribution();
			//
			// int idx = getIndexInEdge(edges(NodeNames.Z), next);
			// GaussianDistribution g = new GaussianDistribution(0, 0);
			// for (int i = 0; i < _zdistributions.size(); i++) {
			// if (i == idx) {
			// continue;
			// }
			// g.setMu(g.getMu() + _zdistributions.get(i).getMu());
			// g.setSigma2(g.getSigma2() + _zdistributions.get(i).getSigma2());
			// }
			// g.setMu(tozDist.getMu() - g.getMu());
			// g.setSigma2(tozDist.getSigma2() - g.getSigma2());
			//
			// m = new Message(g);

			m = trueSkillToZ(next, opt);
		}

		Printer.print("F_ZBR>" + next + ":" + m.getDistribution().toString());

		return m;
	}

	private int _count = 0;
	private double _imp_mean = 0;

	private BeliefMessage trueSkillToZ(GraphNode next, String opt) {
		BeliefMessage m;

		if (_forward) {
			m = edges(NodeNames.RT).nodes().get(0).to(this, opt);

			ArrayList<GaussianDistribution> gd = new ArrayList<GaussianDistribution>();
			for (int i = 0; i < _zdistributions.size(); i++) {
				gd.add(Indicator.trueSkillBackwardEqualSum(i, _zdistributions, _toz,
						(GaussianDistribution) m.getDistribution(), _toR));
			}

			_toz = gd;
			_forward = false;
		}
		int idx = _toz.size() - 1;
		if (next.is(NodeNames.Z)) {
			idx = GraphUtils.getIndexInEdge(edges(NodeNames.Z), next);
		}

		m = new BeliefMessage(_toz.get(idx));
		
		return m;
	}
}
